#!/usr/local/bin/env python

import protocol
import time
import struct
import numpy as np

dl = protocol.Datalink("/dev/ttyUSB0")
channels = ['front', 'reverse', 'left', 'right', 'down']
data = np.zeros(5)              # 5 channels of data
alpha = 1./4                   # For the LPF

# 10 power levels to simulate PWM on the vertical motor
# hand dithered bang-bang.
power = [[0,0,0,0,0,0,0,0,0,0], # 0
         [1,0,0,0,0,0,0,0,0,0], # 1
         [1,0,0,0,0,1,0,0,0,0], # 2
         [1,0,0,1,0,0,1,0,0,0], # 3
         [1,0,1,0,0,1,0,1,0,0], # 4
         [1,0,1,0,1,0,1,0,1,0], # 5
         [1,0,1,0,1,1,1,0,1,0], # 6
         [0,1,1,0,1,1,0,1,1,1], # 7
         [1,1,0,1,1,1,1,0,1,1], # 8
         [1,1,1,1,1,1,1,1,1,1]] # 9

cycle = 0                       # the PWM cycle
target = 350.0                  # target altitude

while True:
    now = time.time()           # For measuring control loop rate

    # This next line decodes the sensor data sent by the blimp into an array
    # in the same order as the 'channels' list declared above.
    new_data = np.array(struct.unpack("hhhhhh", ''.join(dl.receive())[:12])[1:])

    # Apply a low pass filter
    data = alpha*new_data + (1-alpha)*data

    # Calculate the refresh rate
    rate = 1.0/(time.time() - now)

    # Create a dictionary indexed by channel name
    ir = dict(zip(channels, data))

    # Print sensor readings to the screen
    for i in range(len(channels)):
        print '%s %6d  '%(channels[i], data[i]),
    print "H %2.2f"%(rate),

    # Here is a simple P controller
    l = r = 0
    error = ir['down'] - target # positive error when need to go up.
    error /= 5.0                # so a range of 50 gives power.
    sign = 1
    if error < 0:
        sign = -1
    error = int(abs(error))
    if error > 9:               # cap in the range of our PWM settings
        error = 9

    # apply!
    v = sign*power[error][cycle]
    cycle = (cycle+1) % 10

    # Print motor controls
    print 'M %3d %3d %d'%(l,r,v)

    # Send the motor data
    dl.send(struct.pack("bbbb", 1, l,r,v))


